import os
import random
from sklearn.model_selection import train_test_split
import yaml
import shutil

def generate_train_val_files(base_dir, output_dir, skip_folder=None):
    """
    遍历base_dir下的所有子文件夹，为每个子文件夹生成train.txt和val.txt文件。
    然后将这些文件保存到output_dir指定的路径。
    
    Parameters:
    - base_dir: 要遍历的目录路径。
    - output_dir: 输出目录路径。
    - skip_folder: 需要跳过的所有子文件夹及其子文件夹的名称（完全匹配）。
    """
    # 如果输出目录已存在，则删除它并重新创建
    if os.path.exists(output_dir):
        print(f"Deleting existing output directory: {output_dir}")
        shutil.rmtree(output_dir)
    print(f"Creating new output directory: {output_dir}")
    os.makedirs(output_dir, exist_ok=True)

    # 遍历base_dir下的所有子文件夹
    for subdir, dirs, files in os.walk(base_dir):
        # 检查是否需要跳过当前子文件夹
        for dir in dirs:
            subdir_path = os.path.join(subdir, dir)
            # 完全匹配skip_folder
            if skip_folder is not None and os.path.basename(subdir_path) == skip_folder:
                print(f"Skipping directory: {subdir_path}")
                continue

            image_files = [f for f in os.listdir(subdir_path) if f.endswith(('.jpg', '.png'))]
            
            # 如果子文件夹中没有图像文件，跳过
            if not image_files:
                continue

            # 拆分数据集
            train_files, val_files = train_test_split(image_files, test_size=0.2, random_state=42)

            # 生成文件列表
            def write_file_list(file_list, file_name, output_path):
                with open(os.path.join(output_path, file_name), 'w', encoding='utf-8') as f:
                    for file in file_list:
                        abs_path = os.path.abspath(os.path.join(subdir_path, file))
                        f.write(abs_path + '\n')
            
            # 创建子文件夹的名称
            folder_name = os.path.relpath(subdir_path, base_dir).replace('\\', '_').replace('/', '_')

            # 确保子文件夹的输出目录存在
            sub_output_dir = os.path.join(output_dir, folder_name)
            os.makedirs(sub_output_dir, exist_ok=True)

            # 写入训练集和验证集的文件列表
            write_file_list(train_files, 'train.txt', sub_output_dir)
            write_file_list(val_files, 'val.txt', sub_output_dir)

            print(f"{sub_output_dir} - Train set size: {len(train_files)}, Validation set size: {len(val_files)}")

def generate_yaml_file(base_dir, output_yaml_file):
    """
    遍历base_dir下的所有子文件夹，将包含train.txt和val.txt文件的子文件夹路径添加到YAML文件中。
    
    Parameters:
    - base_dir: 要遍历的目录路径。
    - output_yaml_file: 输出YAML文件的路径。
    """
    # 初始化YAML文件内容
    yaml_content = {
        'train': [],
        'val': [],
        'names': {}  # 初始化names字典
    }

    # 遍历base_dir下的所有子文件夹
    for root, dirs, files in os.walk(base_dir):
        # 检查当前目录中是否存在train.txt和val.txt文件
        if 'train.txt' in files and 'val.txt' in files:
            # subdir_path = os.path.relpath(root, base_dir)  # 获取相对路径
            subdir_path = os.path.abspath(root)

            # 添加train.txt文件路径到YAML内容中
            yaml_content['train'].append(os.path.join(subdir_path, 'train.txt'))
            # 添加val.txt文件路径到YAML内容中
            yaml_content['val'].append(os.path.join(subdir_path, 'val.txt'))

    # 写入类别
    yaml_content['names'][0] = 'robot'
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_yaml_file), exist_ok=True)

    # 将YAML内容写入文件
    with open(output_yaml_file, 'w') as yaml_file:
        yaml.dump(yaml_content, yaml_file, default_flow_style=False, sort_keys=False)

    print(f"YAML file has been written to {output_yaml_file}")

if __name__ == "__main__":
    # 划分数据集
    # 数据集路径
    # dataset_dir = '/home/users/xzh/rm_data/test_multi'  # 1. 修改这里
    dataset_dir = r'/home/hxt/nfs_storage/HITCRT/QuanJu/' 
    # 输出目录路径
    # output_dir = os.path.join("/home/users/xzh/ultralytics/rm_train_val/", 'test_multi')  # 2. 修改这里
    output_dir = os.path.join(r"/home/hxt/ultralytics-main/rm_train_val/", 'QuanJu') 
    # 测试集文件夹
    test_dir_name = "val"
    # generate_train_val_files会跳过test_dir_name，便于自己测试模型
    # 注意 output_dir会删除重建，确保更改test_dir_name后不受之前运行结果影响
    generate_train_val_files(dataset_dir,output_dir,test_dir_name)

    # # 路径写入yaml中
    base_directory =  output_dir
    output_yaml_file = r"/home/hxt/ultralytics-main/rm_train_val/RM_data_seg.yaml"
    generate_yaml_file(base_directory, output_yaml_file)


